import torch
import torch.nn as nn
import torch.nn.functional as F

class Generator(nn.Module):
    """Define standard linear + softmax generation step."""

    def __init__(self, d_model, vocab):
        # 初始化函数的输入参数有两个, d_model代表词嵌入维度, vocab.size代表词汇表大小
        super(Generator, self).__init__()
        # 首先就是使用nn中的预定义线性层进行实例化, 得到一个对象self.proj等待使用
        # 这个线性层的参数有两个, 就是初始化函数传进来的两个参数: d_model, vocab_size
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        # 前向逻辑函数中输入是上一层的输出张量x,在函数中, 首先使用上一步得到的self.proj对x进行线性变化,
        # 然后使用F中已经实现的log_softmax进行softmax处理.
        return F.log_softmax(self.proj(x), dim=-1)


if __name__ == "__main__":
    print("=== Generator 生成器测试 ===")

    # 测试参数
    batch_size = 2
    seq_len = 5
    d_model = 512
    vocab_size = 10000

    # 创建生成器和测试数据
    generator = Generator(d_model, vocab_size)
    decoder_output = torch.randn(batch_size, seq_len, d_model)

    print(f"输入shape: {decoder_output.shape}")  # [2, 5, 512]

    # 前向传播
    log_probs = generator(decoder_output)
    print(f"输出shape: {log_probs.shape}")  # [2, 5, 10000]

    # 验证log_softmax性质
    prob_sum = torch.exp(log_probs[0, 0, :]).sum().item()
    print(f"概率和: {prob_sum:.4f} (应该≈1)")

    # 获取预测词汇
    predicted_tokens = torch.argmax(log_probs, dim=-1)
    print(f"预测词汇: {predicted_tokens[0].tolist()}")

    # 测试梯度
    target = torch.randint(0, vocab_size, (batch_size, seq_len))
    loss = F.nll_loss(log_probs.view(-1, vocab_size), target.view(-1))
    loss.backward()
    print(f"损失: {loss.item():.4f}")
    print(f"梯度计算成功: {generator.proj.weight.grad is not None}")

    print("测试完成!")